from time import sleep

import pandas as pd
from MyTT import *
import akshare as ak
import numpy as np
import os
from datetime import datetime
import pymssql
from sqlalchemy import create_engine, distinct, or_, and_
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser

pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db = conf.get(sqlserver, 'db')


# 箱体震荡股票


def filter_stocks():
    sql = "SELECT DISTINCT code FROM STOCK_DAY_DATA"
    r = exec_sql(sql)
    all_stocks = [tok[0] for tok in r]
    # all_stocks = get_stocks()
    ret = list()
    N = 60
    ratio = 20
    rate = 0.8
    max_count = 5

    for i, code in enumerate(all_stocks):
        print(f'len={len(all_stocks)}->i={i}->code={code}')

        res = get_data(code)
        date_list = res['date']
        LOW = res['low']
        HIGH = res['high']
        if len(LOW) < N:
            continue

        A1 = HHV(HIGH, N)
        B1 = LLV(LOW, N)
        factor = (A1 / B1 - 1) * 100
        span = (A1 - B1) * rate
        upper_line = A1 - rate / 2 * span
        lower_line = B1 + rate / 2 * span
        upper_count, lower_count = 0, 0
        max_len = len(LOW)

        for j in range(max_len - 1, max_len):
            if factor[j] < ratio:
                for k in range(N):
                    h = HIGH[j-k]
                    l = LOW[j-k]
                    if h >= upper_line[j]:
                        upper_count += 1
                    if l <= lower_line[j]:
                        lower_count += 1
                    if upper_count >= max_count and lower_count >= max_count:
                        item = [date_list[j], code]
                        ret.append(item)
                        break

        # print(f'code={code}->dt={date_list[-1]}->lower_count={lower_count}->upper_count={upper_count}')
    df = pd.DataFrame(ret, columns=['DATE_T', 'STK_CODE'])
    df.to_excel(r'C:\Users\AndrewX\Desktop\震荡箱体.xlsx', index=False)


def get_stocks():
    stocks = ['000030', '000055', '000078', '000408', '000507', '000561', '000582', '000596', '000601', '000668',
              '000709', '000717', '000788', '000789', '000823', '000836', '000848', '000878', '000921', '000973',
              '000975', '000599', '000407', '000598', '000619', '000623', '000636', '000570']
    stocks = list(set(stocks))
    return stocks


def get_data(code):
    dt = 20230801
    # sql = f"SELECT code, datetime, [close], [open], low, high FROM STOCK_DAY_DATA WHERE code={code} and datetime>={dt}"
    sql = f"SELECT code, datetime, low, high FROM STOCK_DAY_DATA WHERE code={code} and datetime>={dt}"
    r = exec_sql(sql)
    date_ = [tok[1] for tok in r]
    # close = np.array([tok[2] for tok in r])
    # open_ = np.array([tok[3] for tok in r])
    low = np.array([tok[2] for tok in r])
    high = np.array([tok[3] for tok in r])
    # res = {'date': date_, 'close': close, 'open': open_, 'low': low, 'high': high}
    res = {'date': date_, 'low': low, 'high': high}

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


if __name__ == '__main__':
    filter_stocks()
